package oktx.dubbo.interceptor;

import com.alibaba.dubbo.common.Constants;
import com.alibaba.dubbo.common.extension.Activate;
import com.alibaba.dubbo.rpc.*;
import oktx.participant.JtaParticipant;
import oktx.RemoteType;
import oktx.TransactionContext;
import oktx.config.ApplicationContextHolder;
import oktx.dubbo.config.JsonProcessor;
import oktx.config.LocalTxTerminator;
import oktx.thread.TransactionThreadDetach;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.transaction.NotSupportedException;
import javax.transaction.SystemException;
import javax.transaction.Transaction;
import javax.transaction.TransactionManager;

/**
 * Created by dave on 18-5-18 下午9:04.
 * 拦截Dubbo调用，处理请求发起的JTA子事务
 */
@Activate(group = Constants.PROVIDER)
public class DubboProviderInterceptor implements Filter {
    private static final Logger LOG = LoggerFactory.getLogger(DubboProviderInterceptor.class);

    @Override
    public Result invoke(Invoker<?> invoker, Invocation invocation) throws RpcException {
        String json = RpcContext.getContext().getAttachment(JtaParticipant.PRC_TRANSACTION_KEY);
        TransactionContext tc = null;
        if (json != null) { // JTA子事务
            TransactionManager tm = ApplicationContextHolder.getBean(TransactionManager.class);
            if (tm == null) {
                throw new RuntimeException("未找到TransactionManager bean");
            }
            JtaParticipant p = JsonProcessor.parse(json, JtaParticipant.class);

            tc = TransactionContext.getContext(p.getXid());
            if (tc != null) {
                try {
                    tm.resume(tc.getTransaction());
                    TransactionContext.restoreContext(tc);
                } catch (Exception e) {
                    return new RpcResult(e);
                }
                LOG.debug("JTA事务继续：" + tc);
            } else {
                try {
                    tm.setTransactionTimeout(p.getTimeout());
                    tm.begin();
                    Transaction transaction = tm.getTransaction();
                    tc = TransactionContext.newContext(p.getXid(), transaction);
                    tc.setClazz(invoker.getInterface());
                    tc.setMethod(invocation.getMethodName());
                    LOG.debug("JTA事务开启：" + tc);
                } catch (SystemException | NotSupportedException e) {
                    return new RpcResult(e);
                }
            }

            // 将consumer方记入参与者中
            JtaParticipant client = new JtaParticipant();
            client.setXid(p.getXid());
            client.setIp(RpcContext.getContext().getRemoteHost());
            client.setPort(RpcContext.getContext().getRemotePort());
            client.setType(RemoteType.consumer);
            tc.getParticipants().add(client);
        }

        Result r = invoker.invoke(invocation);
        if (r.hasException() && tc != null) {
            LocalTxTerminator terminator = ApplicationContextHolder.getBean(LocalTxTerminator.class);
            try {
                terminator.rollback(tc.getXid(), true);
            } catch (Exception e) {
                return new RpcResult(e);
            }
            LOG.debug("JTA事务本地立即回滚");
        } else if (tc != null) {
            clearThreadLocal(tc);
        }
        return r;
    }

    private void clearThreadLocal(TransactionContext tc) {
        ApplicationContextHolder.getBean(TransactionThreadDetach.class).detach(tc.getTransaction());
        TransactionContext.resetContext();
    }

}
